#gui.py
import os
import tkinter as tk
from tkinter import filedialog
from tkinter import messagebox
import subprocess
from create_folder import FolderCopier  # 导入创建文件的函数
from copy_images import copy_files  # 导入复制文件的函数
from select_yaml import copy_yaml_to_destination  # 导入复制 YAML 文件的函数
import datetime  # 导入日期时间模块

class GUI:
    def __init__(self, master):
        self.master = master
        master.title("ModelAutoTrainer")
        master.geometry("640x640")  # 设置窗口大小

        # 创建主框架
        self.main_frame = tk.Frame(master, padx=20, pady=20)
        self.main_frame.pack(expand=True, fill="both")

        # 新文件夹名称标签
        self.label = tk.Label(self.main_frame, text="输入新文件夹名称:", font=("Arial", 12))
        self.label.pack(anchor="w")

        # 新文件夹名称输入框
        self.entry = tk.Entry(self.main_frame, font=("Arial", 12), width=40)
        self.entry.pack(fill="x", pady=(0, 10))

        # 目标文件夹路径标签
        self.destination_label = tk.Label(self.main_frame, text="", font=("Arial", 12))
        self.destination_label.pack(anchor="w")

        # 源图片文件夹路径标签
        self.source_folder_label = tk.Label(self.main_frame, text="图片路径:", font=("Arial", 12))
        self.source_folder_label.pack(anchor="w")

        # 源图片文件夹路径输入框
        self.source_folder_path = tk.StringVar()
        self.source_folder_entry = tk.Entry(self.main_frame, textvariable=self.source_folder_path, font=("Arial", 12), width=40)
        self.source_folder_entry.pack(fill="x", pady=(0, 10))

        # 选择源图片文件夹路径按钮
        self.source_folder_button = tk.Button(self.main_frame, text="选择图片", command=self.choose_source_folder, font=("Arial", 12))
        self.source_folder_button.pack(pady=(0, 20))

        # 源yaml文件夹路径标签
        self.yaml_file_label = tk.Label(self.main_frame, text="YAML 文件路径:", font=("Arial", 12))
        self.yaml_file_label.pack(anchor="w")

        # 源 YAML 文件路径输入框
        self.yaml_file_path = tk.StringVar()
        self.yaml_file_entry = tk.Entry(self.main_frame, textvariable=self.yaml_file_path, font=("Arial", 12), width=40)
        self.yaml_file_entry.pack(fill="x", pady=(0, 10))

        # 选择 YAML 文件路径按钮
        self.yaml_file_button = tk.Button(self.main_frame, text="选择 YAML 文件", command=self.choose_yaml_file, font=("Arial", 12))
        self.yaml_file_button.pack(pady=(0, 20))

        # 添加 workers、batch_size 和 epochs 输入框
        self.workers_label = tk.Label(self.main_frame, text="Workers:", font=("Arial", 12))
        self.workers_label.pack(anchor="w")
        self.workers_entry = tk.Entry(self.main_frame, font=("Arial", 12), width=10)
        self.workers_entry.pack(fill="x", pady=(0, 10))

        self.batch_size_label = tk.Label(self.main_frame, text="Batch Size:", font=("Arial", 12))
        self.batch_size_label.pack(anchor="w")
        self.batch_size_entry = tk.Entry(self.main_frame, font=("Arial", 12), width=10)
        self.batch_size_entry.pack(fill="x", pady=(0, 10))

        self.epochs_label = tk.Label(self.main_frame, text="Epochs:", font=("Arial", 12))
        self.epochs_label.pack(anchor="w")
        self.epochs_entry = tk.Entry(self.main_frame, font=("Arial", 12), width=10)
        self.epochs_entry.pack(fill="x", pady=(0, 10))

        # 创建按钮框架
        self.button_frame = tk.Frame(self.main_frame)
        self.button_frame.pack(side="bottom", pady=(20, 0), anchor="w")

        # 创建文件夹按钮
        self.create_folder_button = tk.Button(self.button_frame, text="创建文件夹", command=self.create_folder, font=("Arial", 12))
        self.create_folder_button.grid(row=0, column=0, padx=(0, 10))

        # 开始复制图片按钮
        self.copy_images_button = tk.Button(self.button_frame, text="开始复制图片", command=self.copy_images, font=("Arial", 12))
        self.copy_images_button.grid(row=0, column=1, padx=(0, 10))

        # 复制 YAML 文件按钮
        self.copy_yaml_button = tk.Button(self.button_frame, text="复制 YAML 文件", command=self.copy_yaml_file, font=("Arial", 12))
        self.copy_yaml_button.grid(row=0, column=2, padx=(0, 10))  

        # 数据权分按钮
        self.data_rights_button = tk.Button(self.button_frame, text="数据权分", command=self.data_classification, font=("Arial", 12))
        self.data_rights_button.grid(row=0, column=3, padx=(0, 10))

        # 确定按钮
        self.confirm_button = tk.Button(self.button_frame, text="参数确定", command=self.open_parent_directory, font=("Arial", 12))
        self.confirm_button.grid(row=0, column=4)

        # 设置列权重，使输入框扩展
        self.main_frame.pack_propagate(False)

        # 初始化目标文件夹路径
        self.destination_folder = ""

        # 初始化训练参数
        self.workers = None
        self.batch_size = None
        self.epochs = None

        # 初始化日志文件
        self.log_file_path = "log.txt"
        self.log("应用程序启动")

    def log(self, message):
        """记录日志信息到 log.txt 文件"""
        with open(self.log_file_path, "a") as log_file:
            timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            log_file.write(f"{timestamp} - {message}\n")

    def choose_source_folder(self):
        folder_selected = filedialog.askdirectory()
        if folder_selected:
            self.source_folder_path.set(folder_selected)
            self.log(f"选择源图片文件夹: {folder_selected}")

    def create_folder(self):
        new_folder_name = self.entry.get().strip()
        folder_copier = FolderCopier()  # 创建实例

        if not new_folder_name:
            messagebox.showwarning("警告", "请输入新文件夹名称！")
            return

        successful = folder_copier.copy_folder(new_folder_name)  # 使用用户输入的名称

        if successful:
            # 获取目标文件夹路径并更新标签
            self.destination_folder = FolderCopier.destination_folder  # 获取目标文件夹路径
            # 规范化路径，去掉多余部分
            self.destination_folder = os.path.normpath(self.destination_folder)
            self.destination_label.config(text=f"目标文件夹路径: {self.destination_folder}")  # 更新标签文本
            messagebox.showinfo("成功", f"成功复制文件夹 '{new_folder_name}'。")
            self.log(f"成功复制文件夹: {new_folder_name}")
        else:
            messagebox.showwarning("失败", f"复制文件夹 '{new_folder_name}' 失败。")
            self.log(f"复制文件夹失败: {new_folder_name}")

    def choose_yaml_file(self):
        file_selected = filedialog.askopenfilename(filetypes=[("YAML files", "*.yaml")])
        if file_selected:
            self.yaml_file_path.set(file_selected)
            self.log(f"选择 YAML 文件: {file_selected}")

    def copy_yaml_file(self):
        yaml_file_path = self.yaml_file_path.get()  # 获取 YAML 文件路径
        if not self.destination_folder:
            messagebox.showwarning("警告", "请先创建目标文件夹！")
            return

        copy_yaml_to_destination(yaml_file_path, self.destination_folder)  # 调用复制函数
        self.log(f"复制 YAML 文件到目标文件夹: {self.destination_folder}")

    def open_parent_directory(self): 
        # 获取当前工作目录的上一层目录 
        parent_directory = os.path.abspath(os.path.join(os.getcwd(), os.pardir)) 

        # 获取命令参数 
        workers = self.workers_entry.get().strip() 
        batch_size = self.batch_size_entry.get().strip() 
        epochs = self.epochs_entry.get().strip() 
        new_folder_name = self.entry.get().strip() 
        data_file_path = os.path.join('data', new_folder_name, f'{new_folder_name}.yaml') 

        # 构建完整的命令 
        command = f"python train.py --weights yolov5s.pt --data {data_file_path} --workers {workers} --batch-size {batch_size} --epochs {epochs}" 

        # 根据操作系统选择终端命令 
        if os.name == 'nt':  # Windows 
            # 使用 cmd.exe 显示命令
            subprocess.Popen(['start', 'cmd', '/k', f'cd {parent_directory} && echo {command}'], shell=True) 
        else:  # Unix/Linux/Mac 
            # 使用终端显示命令
            subprocess.Popen(['gnome-terminal', '--working-directory', parent_directory, '--', 'bash', '-c', f'echo {command}; exec bash'])  # 适用于 GNOME 终端
            # 或者使用 xterm
            # subprocess.Popen(['xterm', '-e', 'bash', '-c', f'echo {command}; exec bash'])
        
        self.log(f"打开命令行并显示命令: {command}")

    def copy_images(self):
        source_folder = self.source_folder_path.get()
        if not source_folder:
            messagebox.showwarning("警告", "请选择源图片文件夹路径。")
            return
        if not self.destination_folder:
            messagebox.showwarning("警告", "请先创建目标文件夹！")
            return

        # 调用复制文件的函数
        copy_files(source_folder, self.destination_folder)
        self.log(f"复制图片从 {source_folder} 到 {self.destination_folder}")

    def data_classification(self):
        command = [
            'python', 'data_classification.py', self.destination_folder  # 传递目标文件夹路径
        ]

        # 执行命令
        try:
            subprocess.run(command, check=True)
            messagebox.showinfo("成功", "数据权分成功。")
            self.log(f"数据权分成功，目标文件夹: {self.destination_folder}")
        except subprocess.CalledProcessError as e:
            messagebox.showerror("错误", f"数据权分时出错: {e}")
            self.log(f"数据权分时出错: {e}")

if __name__ == "__main__":
    root = tk.Tk()
    app = GUI(root)
    root.mainloop()
#已经完成,操作记录在log.txt中显示